#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
pick_fastest_apt.py
自动测速并设置最快的 APT 源（Ubuntu/Debian）。默认保留 security 源在官方域名，可选一起替换。
仅用标准库：argparse、subprocess、ssl、urllib、time、os、shutil、datetime、json、platform。

评分方法：
- 针对每个镜像，读取该镜像的 InRelease（若不可用则退化为 Release）以及（可选）对应的 *-security/InRelease，
  统计“连接 + 首字节 + 连续读取上限 bytes”的总耗时；多次尝试取最优；两者求和为镜像的分数。
- 分数越小越好。不可达的镜像归为最后。
持久化：
- 生成规范化的 sources.list（包含 main/universe/multiverse 或 main/contrib/non-free 等组件），自动备份旧文件。
- 非 root 时，尝试使用 `sudo tee` 写入；若无 sudo，将把结果打印到标准输出供你手动保存。
- 回滚非常简单：脚本在写入前会把现有 /etc/apt/sources.list 备份为/etc/apt/sources.list.bak-YYYYmmdd-HHMMSS，
  直接恢复即可。默认不替换安全仓库，可以使用--mirror-security参数替换为同一镜像域。
"""

import argparse
import os
import sys
import time
import ssl
import urllib.request
import subprocess
import shutil
from datetime import datetime
from typing import List, Tuple, Optional, Dict

UA = "apt-source-bench/1.0 (+https://www.debian.org/; +https://ubuntu.com/)"
DEFAULT_BYTES = 128 * 1024
DEFAULT_TRIES = 2
DEFAULT_TIMEOUT = 8.0

# === 默认候选镜像（带根路径）===
UBUNTU_CANDIDATES = [
    "https://archive.ubuntu.com/ubuntu",
    "https://mirrors.tuna.tsinghua.edu.cn/ubuntu",
    "https://mirrors.aliyun.com/ubuntu",
    "https://mirrors.ustc.edu.cn/ubuntu",
    "https://mirrors.bfsu.edu.cn/ubuntu",
    "https://mirrors.huaweicloud.com/ubuntu",
    "https://mirrors.cloud.tencent.com/ubuntu",
    "https://mirrors.163.com/ubuntu",
]
DEBIAN_CANDIDATES = [
    "https://deb.debian.org/debian",
    "https://mirrors.tuna.tsinghua.edu.cn/debian",
    "https://mirrors.ustc.edu.cn/debian",
    "https://mirrors.bfsu.edu.cn/debian",
    "https://mirrors.aliyun.com/debian",
]

def detect_distro_codename() -> Tuple[str, Optional[str]]:
    """
    检测发行版（ubuntu/debian）与代号（codename）。
    优先读取 /etc/os-release 与 /etc/lsb-release；退化尝试 lsb_release。
    """
    distro = ""
    codename = None

    # /etc/os-release
    osr = {}
    try:
        with open("/etc/os-release", "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if "=" in line:
                    k, v = line.split("=", 1)
                    osr[k] = v.strip().strip('"')
    except Exception:
        pass

    id_like = (osr.get("ID") or "").lower()
    if "ubuntu" in id_like:
        distro = "ubuntu"
    elif "debian" in id_like:
        distro = "debian"

    codename = osr.get("VERSION_CODENAME") or None

    # /etc/lsb-release 补充
    if codename is None:
        try:
            with open("/etc/lsb-release", "r", encoding="utf-8") as f:
                for line in f:
                    if line.startswith("DISTRIB_CODENAME="):
                        codename = line.split("=", 1)[1].strip()
        except Exception:
            pass

    # lsb_release -sc 作为兜底
    if codename is None:
        try:
            p = subprocess.run(["lsb_release", "-sc"], capture_output=True, text=True, check=False)
            if p.returncode == 0:
                codename = (p.stdout or "").strip()
        except Exception:
            pass

    return distro, codename

def format_secs(s: float) -> str:
    return f"{s*1000:.0f} ms" if s < 1 else f"{s:.2f} s"

def http_read_time(url: str, max_bytes: int, timeout: float) -> Tuple[Optional[float], str]:
    """
    读取最多 max_bytes 字节的耗时（秒）。失败返回 (None, 注记)。
    """
    ctx = ssl.create_default_context()
    start = time.perf_counter()
    first_byte_at = None
    read_bytes = 0
    try:
        req = urllib.request.Request(url, headers={"User-Agent": UA})
        with urllib.request.urlopen(req, timeout=timeout, context=ctx) as resp:
            chunk = resp.read(1)
            if not chunk:
                return None, "empty response"
            first_byte_at = time.perf_counter()
            read_bytes += 1
            while read_bytes < max_bytes:
                c = resp.read(min(8192, max_bytes - read_bytes))
                if not c:
                    break
                read_bytes += len(c)
        elapsed = time.perf_counter() - start
        extra = ""
        if first_byte_at is not None:
            rtt_like = first_byte_at - start
            extra = f" | RTT≈{format_secs(rtt_like)} | read={read_bytes} B"
        return elapsed, f"ok{extra}"
    except Exception as e:
        return None, f"error: {type(e).__name__}: {e}"

def probe_release_urls(base: str, distro: str, codename: str, max_bytes: int, timeout: float,
                       tries: int, include_security_in_score: bool) -> Tuple[str, Optional[float], List[str], str]:
    """
    针对一个镜像 base（含 /ubuntu 或 /debian），构造需要测速的 Release/InRelease 路径并打分。
    返回：(base, 分数 或 None, 实际测试的 URL 列表, 注记)
    """
    tested_urls = []

    # 主仓库
    inrel = f"{base}/dists/{codename}/InRelease"
    rel = f"{base}/dists/{codename}/Release"
    tested_urls.append(inrel)

    # security 仓库（仅当一起替换时纳入评分：Ubuntu 的 {codename}-security；Debian 的 {codename}-security 在 debian-security 命名空间）
    if include_security_in_score:
        if distro == "ubuntu":
            tested_urls.append(f"{base}/dists/{codename}-security/InRelease")
        else:  # debian
            # Debian 的 security 路径通常在 /debian-security
            # 若镜像 base 已是 .../debian，则把 /debian 替换为 /debian-security 以构造对应路径
            if base.endswith("/debian"):
                sec_base = base[:-len("/debian")] + "/debian-security"
            else:
                sec_base = base + "-security"
            tested_urls.append(f"{sec_base}/dists/{codename}-security/InRelease")

    total = 0.0
    notes = []
    for url in tested_urls:
        # InRelease 失败则尝试 Release
        cand = [url]
        if url.endswith("/InRelease"):
            cand.append(url[:-len("InRelease")] + "Release")
        best = None
        last_note = ""
        for _ in range(tries):
            for u in cand:
                t, n = http_read_time(u, max_bytes, timeout)
                last_note = n
                if t is not None:
                    best = t if best is None else min(best, t)
        if best is None:
            return base, None, tested_urls, f"fail at {url} -> {last_note}"
        total += best
        notes.append(f"{url} ok")
    return base, total, tested_urls, "ok"

def rank_mirrors(bases: List[str], distro: str, codename: str, max_bytes: int, timeout: float,
                 tries: int, include_security_in_score: bool):
    results = []
    seen = set()
    uniq = []
    for b in bases:
        b = b.strip().rstrip("/")
        if b and b not in seen:
            uniq.append(b)
            seen.add(b)
    for b in uniq:
        base, score, urls, note = probe_release_urls(b, distro, codename, max_bytes, timeout, tries, include_security_in_score)
        results.append((base, score, urls, note))
    results.sort(key=lambda x: float("inf") if x[1] is None else x[1])
    return results

def make_sources_list(base: str, distro: str, codename: str,
                      include_updates: bool, include_backports: bool, mirror_security_domain: bool) -> str:
    """
    生成标准 sources.list 文本。Ubuntu 与 Debian 的组件命名略不同。
    - Ubuntu：main restricted universe multiverse
    - Debian：main contrib non-free non-free-firmware（自 Debian 12 起）
    security 源：默认保留官方域名；若 mirror_security_domain=True，则与 base 同域。
    """
    lines = []
    if distro == "ubuntu":
        comps = "main restricted universe multiverse"
        # base 用于 main/updates/backports
        lines.append(f"deb {base} {codename} {comps}")
        if include_updates:
            lines.append(f"deb {base} {codename}-updates {comps}")
        if include_backports:
            lines.append(f"deb {base} {codename}-backports {comps}")
        # security
        sec_base = base if mirror_security_domain else "https://security.ubuntu.com/ubuntu"
        lines.append(f"deb {sec_base} {codename}-security {comps}")
    else:  # debian
        comps = "main contrib non-free non-free-firmware"
        # 主仓库
        lines.append(f"deb {base} {codename} {comps}")
        if include_updates:
            lines.append(f"deb {base} {codename}-updates {comps}")
        if include_backports:
            lines.append(f"deb {base} {codename}-backports {comps}")
        # security 仓库（标准路径在 debian-security 命名空间）
        if base.endswith("/debian") and mirror_security_domain:
            sec_base = base[:-len("/debian")] + "/debian-security"
        else:
            sec_base = "https://security.debian.org/debian-security"
        lines.append(f"deb {sec_base} {codename}-security {comps}")
    return "\n".join(lines) + "\n"

def write_sources(content: str, do_apply: bool, do_apt_update: bool) -> Tuple[bool, str]:
    """
    写入 /etc/apt/sources.list，自动备份，并可选执行 apt update。
    - 非 root：尝试使用 `sudo tee`；若无 sudo，则改为打印内容并返回 False 提示手动写入。
    """
    target = "/etc/apt/sources.list"
    ts = datetime.now().strftime("%Y%m%d-%H%M%S")
    backup = f"/etc/apt/sources.list.bak-{ts}"

    if not do_apply:
        # 仅显示将要写入的内容
        return False, f"(dry-run) >>> 将写入 {target} 以下内容：\n" + content

    # 备份
    try:
        if os.path.exists(target):
            shutil.copy2(target, backup)
    except Exception as e:
        return False, f"备份失败：{e}"

    # 写入
    try:
        if os.geteuid() == 0:
            with open(target, "w", encoding="utf-8") as f:
                f.write(content)
            msg = f"已写入 {target}\n备份在：{backup}"
        else:
            # 使用 sudo tee
            p = subprocess.run(["sudo", "tee", target], input=content, text=True, capture_output=True, check=False)
            if p.returncode != 0:
                return False, f"写入失败（sudo tee）：{p.stderr or p.stdout}"
            msg = f"(sudo) 已写入 {target}\n备份在：{backup}"
    except Exception as e:
        return False, f"写入异常：{e}"

    # 可选 apt update
    if do_apt_update:
        try:
            p = subprocess.run(["sudo", "apt-get", "update"], capture_output=True, text=True, check=False)
            msg += "\n\n== apt-get update 输出 ==\n" + (p.stdout or "") + (p.stderr or "")
        except Exception as e:
            msg += f"\napt-get update 失败：{e}"

    return True, msg

def main():
    ap = argparse.ArgumentParser(
        description="自动测速并设置最快 APT 源（Ubuntu/Debian），写入 /etc/apt/sources.list（含自动备份）。"
    )
    ap.add_argument("--distro", choices=["ubuntu", "debian"], default="",
                    help="指定发行版（默认自动检测）。")
    ap.add_argument("--codename", type=str, default="",
                    help="指定代号（默认自动检测），如 ubuntu:noble / debian:bookworm。")
    ap.add_argument("--candidates", type=str, default="",
                    help="自定义候选镜像（逗号分隔，需包含 /ubuntu 或 /debian 根路径）。")
    ap.add_argument("--bytes", type=int, default=DEFAULT_BYTES,
                    help=f"单次读取字节上限（默认：{DEFAULT_BYTES} 字节）。")
    ap.add_argument("--timeout", type=float, default=DEFAULT_TIMEOUT,
                    help=f"连接+读取超时（秒，默认：{DEFAULT_TIMEOUT}）。")
    ap.add_argument("--tries", type=int, default=DEFAULT_TRIES,
                    help=f"每个 URL 的尝试次数（默认：{DEFAULT_TRIES}，取最好的一次）。")
    ap.add_argument("--apply", action="store_true",
                    help="写入 /etc/apt/sources.list（自动备份；需 root 或 sudo）。")
    ap.add_argument("--apt-update", action="store_true",
                    help="写入后执行 `apt-get update`。")
    ap.add_argument("--mirror-security", action="store_true",
                    help="将 security 仓库也替换到同一镜像域（默认保留官方 security 域）。")
    ap.add_argument("--no-updates", action="store_true",
                    help="不写入 -updates 行。")
    ap.add_argument("--no-backports", action="store_true",
                    help="不写入 -backports 行。")
    args = ap.parse_args()

    distro, codename = detect_distro_codename()
    if args.distro:
        distro = args.distro
    if args.codename:
        codename = args.codename
    if not distro or not codename:
        print("无法自动识别发行版或代号。请使用 --distro 与 --codename 指定。", file=sys.stderr)
        sys.exit(2)

    bases = UBUNTU_CANDIDATES if distro == "ubuntu" else DEBIAN_CANDIDATES
    if args.candidates.strip():
        bases += [x.strip() for x in args.candidates.split(",") if x.strip()]

    print(f"== APT 镜像测速开始 ==")
    print(f"发行版：{distro}  |  代号：{codename}  |  读取上限：{args.bytes} B  |  重试：{args.tries}  |  security计入评分：{'是' if args.mirror_security else '否'}")

    results = rank_mirrors(
        bases=bases,
        distro=distro,
        codename=codename,
        max_bytes=args.bytes,
        timeout=args.timeout,
        tries=args.tries,
        include_security_in_score=args.mirror_security
    )

    print("\n== 速度排名（越上越快）==")
    best_base = None
    for i, (base, score, urls, note) in enumerate(results, 1):
        speed = "不可达/失败" if score is None else format_secs(score)
        print(f"{i:02d}. {base:<48} -> {speed}   [{note}]")
        for u in urls:
            print(f"     - {u}")
        if best_base is None and score is not None:
            best_base = base

    if best_base is None:
        print("\n未找到可用镜像。请检查网络、代理（proxy）或更换候选列表。", file=sys.stderr)
        sys.exit(3)

    print(f"\n最快镜像：{best_base}")

    sources = make_sources_list(
        base=best_base,
        distro=distro,
        codename=codename,
        include_updates=not args.no_updates,
        include_backports=not args.no_backports,
        mirror_security_domain=args.mirror_security
    )

    ok, msg = write_sources(sources, do_apply=args.apply, do_apt_update=args.apt_update)
    if not ok:
        print("\n== 将写入的 sources.list 预览 ==")
        print(sources)
        print(msg)
        print("\n你可以手动保存到 /etc/apt/sources.list，然后运行：\n  sudo apt-get update")
        sys.exit(0)
    else:
        print("\n== 写入结果 ==")
        print(msg)

if __name__ == "__main__":
    main()
